import pandas as pd
import numpy as np
import sklearn
from sklearn.model_selection import train_test_split, KFold
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import matplotlib.pylab as pyb

class KNNClassifier:
    def __init__(self, path, selectFeatures, selectTargets, resultPath, isRefresh=False, trainRate=0.7):
        self.path = path
        self.selectFeatures = selectFeatures
        self.selectTargets = selectTargets
        self.resultPath = resultPath
        self.isRefresh = isRefresh
        self.trainRate = trainRate
        self.model = None
        self.scaler = StandardScaler()

    def load_data(self):
        # 读取数据文件
        data = pd.read_csv(self.path)
        # 筛选特征和目标
        X = data.iloc[:, self.selectFeatures].values
        y = data.iloc[:, self.selectTargets].values

        y = y.ravel()

        return X, y

    def shuffle_data(self, X, y):
        if self.isRefresh:
            # 如果需要洗牌，打乱数据
            print("Shuffling data...")
            X, y = sklearn.utils.shuffle(X, y)
        return X, y

    def split_data(self, X, y):
        # 按照 trainRate 拆分数据为训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=self.trainRate, random_state=42)

        print(X_train[0])

        return X_train, X_test, y_train, y_test

    def train(self, X_train, y_train):
        # 对数据进行标准化
        X_train_scaled = self.scaler.fit_transform(X_train)
        # 使用KNN算法训练模型
        self.model = KNeighborsClassifier(n_neighbors=5)  # K值可以根据需要调整
        self.model.fit(X_train_scaled, y_train)

        y_pred = self.model.predict(X_train_scaled)

        print(X_train_scaled[0])

        # 计算准确率
        accuracy = accuracy_score(y_train, y_pred)
        precision = precision_score(y_train, y_pred, average='weighted')
        recall = recall_score(y_train, y_pred, average='weighted')
        f1 = f1_score(y_train, y_pred, average='weighted')
        return {
            "accuracy": format(accuracy, '.3f'),
            "precision": format(precision, '.3f'),
            "recall": format(recall, '.3f'),
            "f1Score": format(f1, '.3f')
        }

    def evaluate(self, X_test, y_test):
        # 对测试数据进行标准化
        X_test_scaled = self.scaler.transform(X_test)
        # 进行预测
        y_pred = self.model.predict(X_test_scaled)

        # 计算准确率
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred, average='weighted')
        recall = recall_score(y_test, y_pred, average='weighted')
        f1 = f1_score(y_test, y_pred, average='weighted')
        # 打印评估结果
        print(f"Accuracy: {accuracy:.4f}")
        print(f"Precision: {precision:.4f}")
        print(f"Recall: {recall:.4f}")
        print(f"F1-Score: {f1:.4f}")
        return {
            "accuracy": format(accuracy, '.3f'),
            "precision": format(precision, '.3f'),
            "recall": format(recall, '.3f'),
            "f1Score": format(f1, '.3f')
        }

    def save_model(self):
        # 保存训练好的模型和数据预处理器（Scaler）
        joblib.dump(self.model, f"{self.resultPath}/model.joblib")
        joblib.dump(self.scaler, f"{self.resultPath}/scaler.joblib")

    def plot_accuracy(self, X_train, y_train, X_test, y_test):
        # 可视化训练和测试的准确率
        accuracies = []
        k_range = range(1, 21)  # 你可以调整K的范围
        for k in k_range:
            model = KNeighborsClassifier(n_neighbors=k)
            model.fit(X_train, y_train)
            y_train_pred = model.predict(X_train)
            y_test_pred = model.predict(X_test)
            train_acc = accuracy_score(y_train, y_train_pred)
            test_acc = accuracy_score(y_test, y_test_pred)
            accuracies.append((k, train_acc, test_acc))

        # 转换为 DataFrame
        accuracy_df = pd.DataFrame(accuracies, columns=["k", "Train Accuracy", "Test Accuracy"])

        # 绘制准确率图表
        plt.figure(figsize=(10, 6))
        sns.lineplot(data=accuracy_df, x="k", y="Train Accuracy", label="Train Accuracy")
        sns.lineplot(data=accuracy_df, x="k", y="Test Accuracy", label="Test Accuracy")
        plt.title("Accuracy vs. k for KNN Classifier")
        plt.xlabel("Number of Neighbors (k)")
        plt.ylabel("Accuracy")
        plt.legend()
        plt.savefig(f"{self.resultPath}/accuracy_plot.png")
        plt.show()

    def cross_validate(self):
        # 加载数据
        X, y = self.load_data()
        # 洗牌数据（如果需要）
        X, y = self.shuffle_data(X, y)

        # 使用 KFold 进行十折交叉验证
        kf = KFold(n_splits=10, shuffle=True)

        fold_results = []

        for train_index, test_index in kf.split(X):
            X_train, X_test = X[train_index], X[test_index]
            y_train, y_test = y[train_index], y[test_index]

            # 训练模型
            train_result = self.train(X_train, y_train)
            # 评估模型
            test_result = self.evaluate(X_test, y_test)
            fold_results.append({
                "trainSet": train_result,
                "testSet": test_result
            })

        # 保存模型
        self.save_model()

        # 计算每个指标的平均值
        avg_results = self.calculate_average(fold_results)

        return avg_results

    def calculate_average(self, fold_results):
        avg_results = {
            "trainSet":{
                "accuracy": 0,
                "precision": 0,
                "recall": 0,
                "f1Score": 0
            },
            "testSet":{
                "accuracy": 0,
                "precision": 0,
                "recall": 0,
                "f1Score": 0
            }
        }

        num_folds = len(fold_results)

        # 对每个指标计算平均值
        for result in fold_results:
            avg_results["testSet"]["accuracy"] += float(result["testSet"]["accuracy"])
            avg_results["testSet"]["precision"] += float(result["testSet"]["precision"])
            avg_results["testSet"]["recall"] += float(result["testSet"]["recall"])
            avg_results["testSet"]["f1Score"] += float(result["testSet"]["f1Score"])

            avg_results["trainSet"]["accuracy"] += float(result["trainSet"]["accuracy"])
            avg_results["trainSet"]["precision"] += float(result["trainSet"]["precision"])
            avg_results["trainSet"]["recall"] += float(result["trainSet"]["recall"])
            avg_results["trainSet"]["f1Score"] += float(result["trainSet"]["f1Score"])

        # 计算平均值
        avg_results["testSet"]["accuracy"] = f"{avg_results['testSet']['accuracy'] / num_folds:.3f}"
        avg_results["testSet"]["precision"] = f"{avg_results['testSet']['precision'] / num_folds:.3f}"
        avg_results["testSet"]["recall"] = f"{avg_results['testSet']['recall'] / num_folds:.3f}"
        avg_results["testSet"]["f1Score"] = f"{avg_results['testSet']['f1Score'] / num_folds:.3f}"

        avg_results["trainSet"]["accuracy"] = f"{avg_results['trainSet']['accuracy'] / num_folds:.3f}"
        avg_results["trainSet"]["precision"] = f"{avg_results['trainSet']['precision'] / num_folds:.3f}"
        avg_results["trainSet"]["recall"] = f"{avg_results['trainSet']['recall'] / num_folds:.3f}"
        avg_results["trainSet"]["f1Score"] = f"{avg_results['trainSet']['f1Score'] / num_folds:.3f}"

        return avg_results

    def run(self):
        # 进行十折交叉验证
        avg_results = self.cross_validate()
        return avg_results

    # def run(self):
    #     # 加载数据
    #     X, y = self.load_data()
    #     # 洗牌数据（如果需要）
    #     X, y = self.shuffle_data(X, y)
    #     # 拆分数据集
    #     X_train, X_test, y_train, y_test = self.split_data(X, y)
    #     # 训练模型
    #     train = self.train(X_train, y_train)
    #     # 评估模型
    #     result = self.evaluate(X_test, y_test)
    #
    #     # 保存模型
    #     self.save_model()
    #
    #     return {
    #         "trainSet": train,
    #         "testSet": result
    #     }


if __name__ == '__main__':
    # 使用示例
    path = "D:\\Code\\WebProjects\\Dataset\\yds-dataset-analyse\\pythonProject\\test\\heart_failure_clinical_records_dataset.csv"
    selectFeatures = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
    selectTargets = [12]
    resultPath = "D:\\Code\\WebProjects\\Dataset\\yds-dataset-analyse\\pythonProject\\test\\result"
    isRefresh = False
    trainRate = 0.8

    knn_classifier = KNNClassifier(path, selectFeatures, selectTargets, resultPath, isRefresh, trainRate)
    result = knn_classifier.run()
    print(result)
